import os
import dotenv
from typing import List

from fastapi import FastAPI
from langchain import hub
from langchain.agents import AgentExecutor
from langchain.agents import create_openai_functions_agent
from langchain.pydantic_v1 import BaseModel, Field
from langchain.tools.retriever import create_retriever_tool
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import FAISS
from langchain_core.messages import BaseMessage
from langchain_openai import ChatOpenAI
from langchain_openai import OpenAIEmbeddings
from langchain_text_splitters import RecursiveCharacterTextSplitter
from langserve import add_routes
from langchain_community.tools import DuckDuckGoSearchRun
from langchain_community.utilities import BingSearchAPIWrapper

dotenv.load_dotenv()
openai_url = os.getenv('OPENAI_API_BASE')
api_key = os.getenv('OPENAI_API_KEY')

# 1. 获取检索器
# 创建一个 WebBaseLoader 对象，加载给定 URL 的网页内容
loader = WebBaseLoader("https://docs.smith.langchain.com/user_guide")
# 载入网页内容
docs = loader.load()

# 初始化 RecursiveCharacterTextSplitter 对象用于文本拆分
text_splitter = RecursiveCharacterTextSplitter()
# 使用文本拆分器将文档分成段落
documents = text_splitter.split_documents(docs)

# 初始化 OpenAIEmbeddings 对象，用于获取文本嵌入
embeddings = OpenAIEmbeddings()
# 从文档中获取嵌入向量并存储
vector = FAISS.from_documents(documents, embeddings)

# 将向量对象转换为检索器
retriever = vector.as_retriever()

# 2. 创建工具
# 检索器工具
retriever_tool = create_retriever_tool(
    retriever,
    "langsmith_search",
    "Search for information about LangSmith. For any questions about LangSmith, you must use this tool!",
)

# 搜索工具
search = DuckDuckGoSearchRun()
tools = [retriever_tool, search]

# 3.创建代理
# 从指定的 Hub 拉取提示模板
prompt = hub.pull("hwchase17/openai-functions-agent")

# 初始化 ChatOpenAI 对象，选择模型为"gpt-3.5-turbo"，设置温度为0
llm = ChatOpenAI(openai_api_key=api_key,openai_api_base=openai_url,model="gpt-3.5-turbo", temperature=0)

# 使用提供的模型、工具和提示创建 OpenAI 函数代理器
agent = create_openai_functions_agent(llm, tools, prompt)

# 初始化 AgentExecutor，传入代理器、工具对象和 verbose 标记为 True
agent_executor = AgentExecutor(agent=agent, tools=tools, verbose=True)

# 4. 应用定义
app = FastAPI(
    title="LangChain Server",
    version="1.0",
    description="A simple API server using LangChain's Runnable interfaces",
)

# 5. 添加路由
class Input(BaseModel):
    # 定义输入 BaseModel 包含字段 input 和 chat_history
    input: str
    chat_history: List[BaseMessage] = Field(
        ...,
        # 为 chat_history 字段添加额外属性，设置 type 为 "chat"，input 为 "location"
        extra={"widget": {"type": "chat", "input": "location"}}
    )


class Output(BaseModel):
    output: str


# 将该配置的agent_executor添加到应用程序app的路由中，路径为 "/agent"
add_routes(
    app,
    # agent_executor配置为使用特定的输入和输出类型
    agent_executor.with_types(input_type=Input, output_type=Output),
    path="/agent",
)

if __name__ == "__main__":
    import uvicorn
    uvicorn.run(app, host="localhost", port=8008)

